"""
高级统计分析功能演示示例

该示例展示了Table类的高级统计分析功能，包括相关系数计算、假设检验和数据分布分析。
"""
from ..dataframe import Table
import matplotlib.pyplot as plt
import scipy.stats as stats
import numpy as np

# 高级统计分析演示
print("===== 高级统计分析功能演示 =====")

# 创建示例数据
data = {
    'age': [25, 30, 35, 40, 45, 50, 55, 60, 65, 70] * 10,
    'income': [50000, 60000, 75000, 90000, 110000, 130000, 150000, 170000, 190000, 210000] * 10,
    'experience': [2, 5, 8, 11, 14, 17, 20, 23, 26, 29] * 10,
    'satisfaction': np.random.randint(1, 10, size=100),
    'department': ['HR', 'Engineering', 'Finance', 'Marketing'] * 25
}

table = Table(data)
print("原始数据前5行:\n", table.head())

# 计算描述性统计
print("\n描述性统计信息:")
print(table.describe())

# 计算相关系数矩阵
corr_matrix = table.corr()
print("\n相关系数矩阵:\n", corr_matrix)

# 绘制相关系数热力图
plt.figure(figsize=(10, 8))
plt.imshow(corr_matrix, cmap='coolwarm', interpolation='none')
plt.colorbar(label='相关系数')
plt.xticks(range(len(corr_matrix.columns)), corr_matrix.columns, rotation=45)
plt.yticks(range(len(corr_matrix.columns)), corr_matrix.columns)
plt.title('特征相关系数热力图')
for i in range(len(corr_matrix.columns)):
    for j in range(len(corr_matrix.columns)):
        plt.text(j, i, f'{corr_matrix.iloc[i, j]:.2f}', ha='center', va='center', color='white')
plt.tight_layout()
plt.show()

# 执行独立样本t检验
engineering_income = table.filter_rows("department == 'Engineering'").get_column('income')
marketing_income = table.filter_rows("department == 'Marketing'").get_column('income')
t_stat, p_value = stats.ttest_ind(engineering_income, marketing_income)
print(f"\n工程部门与市场部门收入独立样本t检验: t={t_stat:.4f}, p={p_value:.4f}")

# 分析年龄与收入的线性关系
slope, intercept, r_value, p_value, std_err = stats.linregress(table.get_column('age'), table.get_column('income'))
print(f"\n年龄与收入线性回归: 斜率={slope:.2f}, 截距={intercept:.2f}, R值={r_value:.4f}, p值={p_value:.4f}")

# 绘制年龄与收入的散点图及回归线
plt.figure(figsize=(10, 6))
plt.scatter(table.get_column('age'), table.get_column('income'), alpha=0.5, label='原始数据')
plt.plot(table.get_column('age'), intercept + slope * table.get_column('age'), 'r', label=f'回归线 (R={r_value:.4f})')
plt.title('年龄与收入关系散点图及回归线')
plt.xlabel('年龄')
plt.ylabel('收入')
plt.legend()
plt.tight_layout()
plt.show()

# 分析满意度得分分布
plt.figure(figsize=(10, 6))
plt.hist(table.get_column('satisfaction'), bins=9, edgecolor='black', alpha=0.7)
plt.title('员工满意度得分分布')
plt.xlabel('满意度得分')
plt.ylabel('员工数量')
plt.tight_layout()
plt.show()